import csv
import os
import glob
from sphfile import SPHFile

# 输入文件夹路径和文件名（不含扩展名）
sph_folder = '/root/autodl-tmp/TED/TEDLIUM_release1/train/sph'
csv_folder = '/root/autodl-tmp/TEDspliter/train/stm'
output_folder = '/root/autodl-tmp/TEDspliter/train/sph'

# 确保输出文件夹存在
os.makedirs(output_folder, exist_ok=True)

# 获取所有 SPH 文件
sph_files = glob.glob(os.path.join(sph_folder, '*.sph'))

# 对每一个SPH文件
for sph_file in sph_files:
    # 获取输入文件的基本名称（无路径）
    base_name = os.path.basename(sph_file)
    # 创建输出文件夹和csv文件的路径
    specific_output_folder = os.path.join(output_folder, os.path.splitext(base_name)[0])
    os.makedirs(specific_output_folder, exist_ok=True)
    csv_file = os.path.join(csv_folder, os.path.splitext(base_name)[0] + '.csv')
    with open(csv_file,'r') as csv_file_entity:
        # audio = AudioSegment.from_file(sph_file, format="sph")
        audio = SPHFile(sph_file)
        reader = csv.DictReader(csv_file_entity)
        for i, row in enumerate(reader):
            # 计算切割的开始和结束时间（秒）
            start_time = float(row['xmin'])
            end_time = float(row['xmax'])
            
            # 切割音频片段
            audio.write_wav(os.path.join(specific_output_folder,f"segment_{i}.wav"), start_time, end_time)

    print(f"Audio segments have been saved to '{specific_output_folder}'")

